// Copyright 2023 Specter Ops, Inc.
//
// Licensed under the Apache License, Version 2.0
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

//go:build ignored
// +build ignored

package main

// This program generates several files for enumerating the different media type definitions.
// It can be invoked by running go generate

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"go/format"
	"io"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"text/template"
	"time"

	"github.com/specterops/bloodhound/slicesext"
)

var packageTemplate = template.Must(template.New("").Parse(`// Code generated by gen.go; DO NOT EDIT.
// This file was generated at {{ .Timestamp }}
// using data from {{ .URL }}
package mediatypes

const (
{{- range .MediaTypes }}
	{{ printf "%s%s MediaType = \"%s\" // %s" .Prefix .Symbol .Template .Reference }},
{{- end }}
)
`))

type TemplateData struct {
	Timestamp  time.Time
	URL        string
	MediaTypes []MediaType
}

type MediaType struct {
	Prefix    string
	Symbol    string
	Template  string
	Reference string
}

func main() {
	categories := []string{
		"application",
		"audio",
		"font",
		"image",
		"message",
		"model",
		"multipart",
		"text",
		"video",
	}

	for _, category := range categories {
		var (
			url              = fmt.Sprintf("https://www.iana.org/assignments/media-types/%s.csv", category)
			categoryFilename = category + ".go"
		)

		if _, err := os.Lstat(categoryFilename); err == nil {
			// Only generate the file if it has been removed as generation requires a call to IANA
			continue
		}

		if res, err := http.Get(url); err != nil {
			log.Fatal(err)
		} else if mediaTypes, err := mediaTypesFromCSV(category, res.Body); err != nil {
			log.Fatal(err)
		} else if buf, err := fillTemplate(url, mediaTypes); err != nil {
			log.Fatal(err)
		} else if formatted, err := format.Source(buf.Bytes()); err != nil {
			log.Fatal(err)
		} else if out, err := os.Create(categoryFilename); err != nil {
			log.Fatal(err)
		} else {
			defer out.Close()
			if _, err := out.Write(formatted); err != nil {
				log.Fatal(err)
			}
		}
	}
}

func fillTemplate(url string, mediatypes []MediaType) (bytes.Buffer, error) {
	var buf bytes.Buffer
	return buf, packageTemplate.Execute(&buf, TemplateData{
		Timestamp:  time.Now(),
		URL:        url,
		MediaTypes: mediatypes,
	})
}

func mediaTypesFromCSV(category string, readcloser io.ReadCloser) ([]MediaType, error) {
	defer readcloser.Close()
	if regex, err := regexp.Compile(`\W`); err != nil {
		return nil, err
	} else if types, err := csv.NewReader(readcloser).ReadAll(); err != nil {
		return nil, err
	} else {
		entries := []MediaType{}
		for _, row := range types[1:] {
			if row[1] != "" {
				symbol := strings.Join(slicesext.Map(regex.Split(row[0], -1), func(word string) string {
					return strings.Title(word)
				}), "")
				entries = append(entries, MediaType{
					Prefix:    strings.Title(category),
					Symbol:    symbol,
					Template:  row[1],
					Reference: row[2],
				})
			}
		}
		return entries, nil
	}
}
